clc;
close all;
clear;
warning off
addpath('C:/casadi-windows-matlabR2016a-v3.5.5')
import casadi.*
%0 1腿号
%
%2 3
%------------------行为规划--------------------
Sect=4;% 相序数量
N = 40; %    总分段比例  多相序为增加MPC控制精度应该  计算出个阶段末端状态作为约束  分段进行MPC控制  必须是偶数  
T = 0.8; %   总时间长度 可以改成按时间规划相序
dt_val = repmat(T/(N),1,N)';%每个相序的分段比例

gen_lib= 1;
Fmax=      128;
max_jump_z=0.8;%最高跳跃高度  但是不能保证一定跳跃到
min_damp_z=0.15;%最低限制高度
max_lift_spdz=3.5;%最大离地速度
z_init=0.2;%起始站立高度
world.mu=0.5;%摩擦系数

q_init_val = [0 0 0 ...
              0 0 z_init]';%初始状态 rpy  xyz
qd_init_val = [0 0 0 0 0 0]';

q_term_ref  = [1*pi*0   -22/57.3*0      pi*0  ...
               0.4+1e-5       0          0.1 ]';%终端位置 rpy  xyz
qd_term_ref = [0 0 0, 0 0 0]';

%Phase_duty=[int16(N/Sect*0.2) int16(N/Sect*0.3) int16(N/Sect*0.49) int16(N/Sect*0.01)];%相序百分比 下蹲 蹬腿  腾空  缓冲
cs_val = [repmat([1 1 1 1]', 1, 15)   repmat([1 1 1 1]', 1, 5)    repmat([0 0 0 0]', 1, 19)  repmat([0 0 0 0]', 1, 1)]';%MPCtable相序  
%接触状态
%cs_TD_val = zeros(Sect,N-1)';
%% %状态的权重
weight.QX = [10 10 10, 10 10 10, 10 10 10, 10 10 10 ]';%系统状态代价的权重 轨迹
weight.QN = [10 10 10, 50 50 150, 10 10 10, 10 10 10 ]';%终端代价
weight.Qc = [0.001 0.001 0.001]';
weight.Qf = [0.0001 0.0001 0.001]';
%%  %物理参数

Body.m = 5+0.25*8;%机器人质量
%机身惯量
I=[0.059150 0.101150 0.046240]*4;
Body.Ib = [I(1)   0      0;%roll
           0      I(2)    0;%pit
           0      0      I(3) ];%yaw 转动惯量矩阵
Body.length_body=0.34;
Body.width_body=0.26;
%0 1腿号
%
%2 3
Body.hipPos=[Body.length_body/2, Body.length_body/2,  -Body.length_body/2, -Body.length_body/2;%髋关节位置
             Body.width_body/2, -Body.width_body/2,    Body.width_body/2,  -Body.width_body/2;
             0,    0,    0, 0];

     endPos=[0,         0,          0,          0;%足端跨关节位置  %初始化足端位置
             0,         0,          0,          0;
             -z_init,   -z_init,   -z_init,    -z_init];
world.g = 9.8;%重力加速度

%%  构造微分方程

Xk=SX.sym('Xk', 12, 1);%cassis 符号类
n_state=size(Xk,1);
Fk=SX.sym('Uk', 12, 1);
n_F=size(Fk,1);
Rk=SX.sym('Rk', 12, 1);
n_r=size(Rk,1);
%%  计算微分方程
I3=eye(3);%单位矩阵
Rbody=rotsb(Xk(1:3));
cy = cos(Xk(3));
sy = sin(Xk(3));
cp = cos(Xk(2));
sp = sin(Xk(2));

R_yaw =[cy sy 0;
        -sy cy 0;
        0 0 1];%世界到机身
R_w=[cy/cp,sy/cp,0;
    -sy,cy,0;
    cy*sp/cp,sy*sp/cp,1];
Ig = Rbody*Body.Ib*Rbody';
Ig_inv=Ig\I3;

%单刚体动力学模型
A = [zeros(3) zeros(3) R_yaw zeros(3)  ;
     zeros(3) zeros(3) zeros(3) I3 ;
     zeros(3) zeros(3) zeros(3) zeros(3);
     zeros(3) zeros(3) zeros(3) zeros(3) ;
    ];%状态矩阵

% AA=A;
% AA(1:3,7:9)=R_w;
B=[ zeros(3)           zeros(3)           zeros(3)            zeros(3);
    zeros(3)           zeros(3)           zeros(3)            zeros(3);
    Ig_inv*Skew(Rk(1:3)) Ig_inv*Skew(Rk(4:6)) Ig_inv*Skew(Rk(7:9))  Ig_inv*Skew(Rk(10:12));
    I3/Body.m   I3/Body.m   I3/Body.m    I3/Body.m;];%控制矩阵
g=zeros(12,1);
g(12)=-world.g;%扩展一维度重力加速度
dotX=A*Xk+B*Fk+g;%构造微分动力学的符号方程

%%  定义函数
f=Function('f',{Xk;Fk;Rk},{dotX},{'input_states','control_inputs','foot_input'},{'dotX'});

% X_init = [0;0.0;0; 0.0;0.0;0.5 ;0;0;0; 0;0;0;-9.8];%初始状态变量
% f(X_init,zeros(12,1),zeros(12,1))%测试函数正常否

%%  构造代价和约束 变量定义
X = SX.sym('X', n_state, N+1); % N+1步状态
F = SX.sym('F', n_F, N); % N步内的控制 力控制
r = SX.sym('r', n_r, N); % N步内的控制 足端位置控制

RefX = SX.sym('RefX', n_state, N+1); % N步内的控制输出
RefF = SX.sym('RefF', n_F, N); % N步内的控制输出
Refr = SX.sym('Refr', n_r, N); % N步内的控制输出
ContactState=SX.sym('ConState', 4, N);
obj=0;
%%  构造代价和约束 变量定义
mu_inv = 1.0/world.mu;
%摩擦约束
f_block =[ mu_inv, 0,  -1.0;
          -mu_inv, 0,  -1.0;
           0,  mu_inv, -1.0;
           0, -mu_inv, -1.0;];

kin_box_x = 0.15;%运动学约束
kin_box_y = 0.15;
kin_box_z = 0.3;%腿最长

Kin_block =[ 1, 0,  0,-kin_box_x;%髋关节坐标系
            -1, 0,  0,-kin_box_x;
             0, 1,  0,-kin_box_y;
             0, -1, 0,-kin_box_y;
             0, 0,  1,-min_damp_z;%腿最短
             0, 0, -1,-kin_box_z];

endPos_Body=Body.hipPos+endPos;
Phip_swing=reshape(endPos_Body,[],1);
%%  约束暂存变量定义 %初态约束
%初态约束
defect_init=RefX(:,1)-X(:,1);%12*1 eq

defect_state=SX.zeros(12*(N+1),1);%12(N+1)*1 eq
defect_FootOnGround=SX.zeros(4*(N),1);%4(N)*1 eq
defect_footStance=SX.zeros(12*(N),1);%(3*4)(N)*1 eq
n_equa_c=12+12*(N+1)+4*(N)+12*(N);
%共
defect_legLimits=SX.zeros(24*(N),1);%(4*6)(N)*1
defect_footforce=SX.zeros(16*(N),1);%(4*4)(N)*1 xy摩擦约束4个
defect_ForceNormal=SX.zeros(N,1);% (N)*1
defect_footswing=SX.zeros(4*(N),1);%4(N)*1
n_inequa_c=24*(N)+16*(N)+N+4*(N);
%%	约束和代价计算
for k = 1:N     
    %%	代价计算    
    Xk=X(:,k);
    Fk=F(:,k);
    rk=r(:,k);
    Pk=repmat(Xk(4:6),4,1)+rk;
    ContactStatek=ContactState(:,k);
    dtk=dt_val(k);
    
    X_err = Xk - RefX(:,k);                                         % 基座状态误差
    pf_err = repmat(Xk(4:6),4,1) + Phip_swing - Pk;                 % 悬空时约束foot位置 不影响跳跃  只影响最终足端状态输出的结果
    U_err = Fk - RefF(:,k);                                         % GRF 误差
    obj = obj + (X_err'*diag(weight.QX)*X_err+...                   % 误差求和
          pf_err'*diag(repmat(weight.Qc,4,1))*pf_err+...
          U_err'*diag(repmat(weight.Qf,4,1))*U_err)*dtk;
    %% 约束计算
    %状态约束
    %runge kutta method
%     k1 = f(Xk,Fk,Pk);   % new
%     k2 = f(Xk + dtk/2*k1,Fk,Pk); % new
%     k3 = f(Xk + dtk/2*k2,Fk,Pk); % new
%     k4 = f(Xk + dtk*k3,Fk,Pk); % new
%     st_next_RK4=Xk +dtk/6*(k1+2*k2+2*k3+k4); % new
%     defect_state((k-1)*12+1:(k-1)*12+12)=X(:,k+1)-(st_next_RK4);
     defect_state((k-1)*12+1:(k-1)*12+12)=X(:,k+1)-(Xk+f(Xk,Fk,rk)*dtk);

    
    %法向力大于0 不等式
    defect_ForceNormal(k)=-dot(Fk,repmat([0;0;1],4,1));
    %结合法向力大于0，摩擦约束来约束摆动中力为0 和最大力 不等式
    defect_footswing((k-1)*4+1:(k-1)*4+4)=Fk([3,6,9,12])-ContactStatek.*repmat(1000,4,1);
    for leg=1:4
        xyz_idx = 3*(leg-1)+1:3*(leg-1)+3;
        %脚在地上约束 0是此时地面高度等式
        defect_FootOnGround((k-1)*4+leg)=ContactStatek(leg)*Pk(3*(leg-1)+3);
        %限制腿长 限制范围不等式
        Rbody=rotsb(Xk(1:3));
        endWorld=Rbody*endPos_Body+Xk(4:6);%全局足端位置
        p_rel = (Pk(xyz_idx) - endWorld(:,leg));%hip->足端  足端矢量
        defect_legLimits((k-1)*24+(leg-1)*6+1:(k-1)*24+(leg-1)*6+6)= Kin_block*[p_rel;1];%运动学约束
        %接触中脚不滑动
        if (k < N)
            Pk1=repmat(X(4:6,k+1),4,1)+r(:,k+1);
            defect_footStance((k-1)*12+(leg-1)*3+1:(k-1)*12+(leg-1)*3+3)=ContactStatek(leg)*(Pk1(xyz_idx)-Pk(xyz_idx));
        end
        %摩擦约束 不等式
        defect_footforce((k-1)*16+(leg-1)*4+1:(k-1)*16+(leg-1)*4+4)=f_block*Fk(xyz_idx);
    end
end
%%	约束生成
g=[defect_init;defect_state;defect_FootOnGround;defect_footStance;...
    defect_legLimits;defect_footforce;defect_ForceNormal;defect_footswing];
display_str=['等式约束数量',num2str(n_equa_c),'   不等式约束数量',num2str(n_inequa_c)];
disp(display_str);
%%	终端 cost
X_err = X(:,end)-RefX(:,end);    % 终端 cost
obj = obj + X_err'*diag(weight.QN)*X_err;


%%	构造问题和问题变量
OPT_variables = [reshape(X,n_state*(N+1),1);reshape(F,n_F*N,1);reshape(r,n_r*N,1)];
OPT_Param = [reshape(RefX,n_state*(N+1),1);reshape(RefF,n_F*N,1);reshape(Refr,n_r*N,1);reshape(ContactState,4*N,1)];
nlp_prob =struct('f', obj, 'x',OPT_variables,'p',OPT_Param, 'g',g );
%%  优化设置
opts_setting.expand =true;
opts_setting.ipopt.max_iter=1500;
opts_setting.ipopt.print_level=0;
opts_setting.ipopt.acceptable_tol=1e-4;
opts_setting.ipopt.acceptable_obj_change_tol=1e-6;
opts_setting.ipopt.tol=1e-4;
opts_setting.ipopt.nlp_scaling_method='gradient-based';
opts_setting.ipopt.constr_viol_tol=1e-3;
opts_setting.ipopt.fixed_variable_treatment='relax_bounds';


%% 构造求解器
if gen_lib
solver = nlpsol('solver', 'ipopt', nlp_prob,opts_setting);%可以在线修改足端位置
else
solver = casadi.nlpsol('solver', 'ipopt', ['nlp.so'],opts_setting);%不需要nlp_prob
end
% tic;
% [status1,cmdout] = system(command,'-echo'); % compile the file % takes a few minutes
% t_comp=toc;
% fprintf('Done compiling in :%.2f s\n',t_comp)
%%	约束上下界面 args
args.lbg(1:n_equa_c) = 0;  % -1e-20  % Equality constraints
args.ubg(1:n_equa_c) = 0;  % 1e-20   % Equality constraints

args.lbg(n_equa_c+1 : n_equa_c+ n_inequa_c) = -inf; % inequality constraints
args.ubg(n_equa_c+1 : n_equa_c+ n_inequa_c) = 0; % inequality constraints
%%	决策变量上下界面 args
%%  状态上边界
tempub=[Body.m*world.g*world.mu*6; Body.m*world.g*world.mu*6 ;Fmax];
args.ubx=[];
UBx=[pi*3*ones(3,1);10*ones(2,1);1;...
     pi*3*ones(3,1);40*ones(3,1)];%状态上届 约束跳的最高高度
UBx(6)=max_jump_z;
UBx(12)=max_lift_spdz;
UBx=repmat(UBx,N+1,1);
UBf=[tempub;tempub;tempub;tempub];
UBf=repmat(UBf,N,1);
UBp=repmat([0.4;0.4;inf],4,1);
UBp=repmat(UBp,N,1);
args.ubx=[args.ubx;UBx;UBf;UBp];
%%  状态下边界
templb=[-Body.m*world.g*world.mu*6; -Body.m*world.g*world.mu*6 ;0];%力状态
args.lbx=[];
LBx=[-pi*3*ones(3,1);-10*ones(2,1);min_damp_z;...
     -pi*3*ones(3,1);-40*ones(3,1)];%状态下界
LBx=repmat(LBx,N+1,1);
LBf=[templb;templb;templb;templb];
LBf=repmat(LBf,N,1);
LBp=repmat([-0.4;-0.4;-inf],4,1);
LBp=repmat(LBp,N,1);
args.lbx=[args.lbx;LBx;LBf;LBp];
%%
 
% c_init_val = repmat(q_init_val(4:6),4,1)+...
%     diag([1 1 1, 1 -1 1, -1 1 1, -1 -1 1])*repmat([0.2 0.1 -q_init_val(6)],1,4)';

c_ref = diag([1 1 1, 1 -1 1, -1 1 1, -1 -1 1])*repmat([0.2 0.1 -z_init],1,4)';%初始化足端位置  --------------------足端位置
f_ref = zeros(12,1);

%% set parameter values 设定期望运动轨迹
for i = 1:6%对状态线性插值
    Xref_val(i,:)   = linspace(q_init_val(i),q_term_ref(i),N+1);%决定轨迹末端位置
    Xref_val(6+i,:) = linspace(qd_init_val(i),qd_term_ref(i),N+1);
end
% Z向抛物线
a=[Xref_val(4,1),Xref_val(4,N/2),Xref_val(4,N)];%x
b=[q_init_val(6),q_term_ref(6),q_init_val(6)+0.0];%z
Xref_val(6,:) =interp1(a,b,Xref_val(4,:),'spline'); %高度方向做Spline插值
Uref_val=zeros(24,N);
r_ref=zeros(12,N);
for leg = 1:4
    for xyz = 1:3
        Uref_val(3*(leg-1)+xyz,:)= Xref_val(xyz+3,1:end-1) +c_ref(3*(leg-1)+xyz);%F 
        r_ref(3*(leg-1)+xyz,:)= c_ref(3*(leg-1)+xyz);%
        Uref_val(12+3*(leg-1)+xyz,:) = f_ref(xyz).*ones(1,N);%P
    end
end

if(1)%线性插值
    for i = 1:6
        Xref_val(i,:)   = linspace(q_init_val(i),q_term_ref(i),N+1);
        Xref_val(6+i,:) = linspace(qd_init_val(i),qd_term_ref(i),N+1);
    end
    for leg = 1:4
        for xyz = 1:3
            Uref_val(3*(leg-1)+xyz,:)    = Xref_val(xyz,1:end-1) + c_ref(3*(leg-1)+xyz);
            Uref_val(12+3*(leg-1)+xyz,:) = f_ref(xyz).*ones(1,N);
        end
    end
end
F_ref=Uref_val(13:24,:);

args.p=[reshape(Xref_val,n_state*(N+1),1);reshape(F_ref,n_F*N,1);reshape(r_ref,n_r*N,1);reshape(cs_val',4*N,1)];%送入了轨迹约束 相序约束
args.x0=[reshape(Xref_val,n_state*(N+1),1);reshape(F_ref,n_F*N,1);reshape(r_ref,n_r*N,1)];%系统初始状态
%-----------求解
sol=solver('x0',args.x0,'lbx', args.lbx,'ubx', args.ubx,'lbg', args.lbg,'ubg', args.ubg,'p',args.p);%调用求解器  输入数据

x_li=sol.x(1:n_state*(N+1));
x_li=reshape(full(x_li),n_state,(N+1));%质心世界下的位置

f_li=sol.x(n_state*(N+1)+1:n_state*(N+1)+n_F*N);
f_li=reshape(full(f_li),n_F,N);%足端力

r_li=sol.x(n_state*(N+1)+n_F*N+1:n_state*(N+1)+n_F*N+n_r*N);
r_li=reshape(full(r_li),n_F,N);%机体下足端位置
p_li=r_li+repmat(x_li(4:6,1:end-1),4,1);%世界足端位置

%------------导出
if gen_lib
    solver = nlpsol('solver','ipopt',nlp_prob,opts_setting);
    disp('Solver without Simple Bounds generated');
    c_file_name = 'nlp';
    disp('Generating c code');
    opts = struct('main', false,...
    'mex', true);
    solver.generate_dependencies([c_file_name,'.c'],opts);% generete helper functions
    disp('Done generating .c file');
    if 1
        disp('Compiling .c file (takes ~15 minutes)');
        command = ['gcc -fPIC -shared -O0 ', c_file_name,'.c -o ',c_file_name,'.so'];  %00-35s o1-229.17
        tic;
        [status1,cmdout] = system(command,'-echo'); % compile the file % takes a few minutes
        t_comp=toc;
        fprintf('Done compiling in :%.2f s\n',t_comp)
    end
end
%------------------------记录
writematrix(x_li,'x_li.csv');
writematrix(f_li,'f_li.csv');
%------------------------绘图
figure(1);
subplot(2,1,1)
plot(x_li(6,:));grid on;%绘制质心运行轨迹
subplot(2,1,2)
plot(x_li(12,:));grid on;%绘制质心运行轨迹   RPY XYZ  DRPY DXYZ


figure(2);
subplot(4,1,1)
plot(f_li(3,:));%绘制足端力
hold on; grid on;
subplot(4,1,2)
plot(f_li(6,:));%绘制足端力
hold on; grid on;
subplot(4,1,3)
plot(f_li(9,:));%绘制足端力
hold on; grid on;
subplot(4,1,4)
plot(f_li(12,:));%绘制足端力
grid on;
 
figure(5);
subplot(4,1,1)
plot(f_li(1,:));%绘制足端力
hold on; grid on;
subplot(4,1,2)
plot(f_li(4,:));%绘制足端力
hold on; grid on;
subplot(4,1,3)
plot(f_li(7,:));%绘制足端力
hold on; grid on;
subplot(4,1,4)
plot(f_li(10,:));%绘制足端力
grid on;

figure(3);
pic_num = 1;%保存gif用
time=['NLP','_',datestr(datetime('now'),'yyyy-mm-dd-HH-MM'),'_Animated.gif'];
for i=1:N
    cube_animate(x_li(:,i),i,p_li(:,i),~cs_val(i,:),[0;0;0;0],...
        f_li(:,i),3,[],[],[],[],[],[-20,14],dt_val,[]);
pause(0.01);%影响绘画
end
%% 工具函数
function rotxm=rotx(theta)
s=sin(theta);
c=cos(theta);
% rotxm=[1,0,0;
%     0,c,s
%     0,-s c]';
rotxm=[1,0,0;
    0,c,-s
    0,s c];
end

function rotym=roty(theta)
s=sin(theta);
c=cos(theta);
% rotym =[c,0,-s;
%     0,1,0;
%     s,0,c]';
rotym =[c,0,s;
    0,1,0;
    -s,0,c];
end

function rotzm=rotz(theta)
s=sin(theta);
c=cos(theta);

% rotzm=[c,s,0;
%     -s,c,0;
%     0,0,1]';
rotzm=[c,-s,0;
    s,c,0;
    0,0,1];
end
%Rsb
function R=rotsb(theta)%构造旋转矩阵
% R=rotx(theta(1))*roty(theta(2))*rotz(theta(3));
R=rotz(theta(3))*roty(theta(2))*rotx(theta(1));

end

function s=Skew(in)
s = [0 -in(3) in(2);
    in(3) 0 -in(1);
    -in(2) in(1) 0];
end